package org.ne.springboot.utils.PaperHandleUtils.PaperContentCheckUtils;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import okhttp3.*;
import okhttp3.EventListener;
import org.jetbrains.annotations.NotNull;

import java.util.*;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicReference;

import com.google.gson.Gson;
import org.ne.springboot.pojo.PromptProcessor;
import org.ne.springboot.utils.YamlUtils;

import javax.crypto.Mac;
import javax.crypto.spec.SecretKeySpec;

import java.net.URL;
import java.nio.charset.StandardCharsets;
import java.text.SimpleDateFormat;


/**
 * @author wey
 */


public class SparkClient {


    /*  hostUrl: https://spark-api.xf-yun.com/v3.5/chat
  domain: generalv3.5
  appid: 6f84c7f7
  apiKey: 129b3fa51b00e89eb642ca505113e40e
  apiSecret: YTU2NTgwZTJkYzk3MjAzYWNkYWQxZGE4
    *
    * */
    private Map<String, Object> map = new HashMap<>();
    private Map<String, Object> format = new HashMap<>();
    private String prompt = "default";
    private String hostUrl;
    private String domain;
    private String appid ;
    private String apiSecret ;
    private String apiKey ;
    private static final PromptProcessor PROMPT_PROCESSOR = new PromptProcessor();

    private static final OkHttpClient CLIENT = new OkHttpClient.Builder()
            .eventListener(new EventListener() {
                @Override
                public void connectionReleased(@NotNull Call call, @NotNull Connection connection) {

                    System.out.println("连接释放: " + connection.socket());

                }
            })
            .build();
    private static final ExecutorService EXECUTOR = new ThreadPoolExecutor(
        5,
        // 核心线程数
        20,
        // 最大线程数
        60L, TimeUnit.SECONDS,
        // 空闲线程存活时间
        new LinkedBlockingQueue<>(100),
        // 任务队列容量
        Executors.defaultThreadFactory(),
        // 线程工厂
        new ThreadPoolExecutor.CallerRunsPolicy()
        // 拒绝策略：由调用线程直接执行
    );


    public Map<String, Object> processRequest(String question, String content) {
        hostUrl = (String) map.get("hostUrl");
        domain = (String) map.get("domain");
        appid = (String) map.get("appid");
        apiSecret = (String) map.get("apiSecret");
        apiKey = (String) map.get("apiKey");

        Map<String, Object> result = new HashMap<>();
        String sessionId = UUID.randomUUID().toString();

        try {

            String combinedInput = PROMPT_PROCESSOR.buildPrompt(prompt, question, content, format);

            // 创建超时控制的CompletableFuture
            CompletableFuture<String> responseFuture = CompletableFuture.supplyAsync(() -> {
                try {
                    return sendToSparkAPI(sessionId, combinedInput);
                } catch (Exception e) {
                    throw new CompletionException("API请求失败: " + e.getMessage(), e);
                }
            }, EXECUTOR);

            // 异步处理响应
            CompletableFuture<Map<String, Object>> processingFuture = responseFuture.thenApply(answer -> {
                try {
                    // 提取JSON部分


                    int start = answer.indexOf("```json") + "```json".length();
                    int end = answer.lastIndexOf("```");
                    String jsonStr = answer.substring(start, end).trim();

                    // 解析JSON
                    JSONObject jsonObject = JSON.parseObject(jsonStr);
                    Map jsonMap = JSONObject.toJavaObject(jsonObject, Map.class);

                    // 构建成功结果
                    Map<String, Object> successResult = new HashMap<>();
                    successResult.put("status", "success");
                    successResult.put("answer", jsonMap);
                    successResult.put("sessionId", sessionId);

                    return successResult;
                } catch (Exception e) {
                    throw new CompletionException("响应解析失败: " + e.getMessage(), e);
                }
            });

            // 等待处理完成（带超时控制）
            result = processingFuture.get(30, TimeUnit.SECONDS);

        } catch (TimeoutException e) {
            result.put("status", "error");
            result.put("message", "处理超时，请重试");
        } catch (ExecutionException e) {
            result.put("status", "error");
            result.put("message", "执行错误: " + e.getCause().getMessage());
        } catch (Exception e) {
            result.put("status", "error");
            result.put("message", "未知错误: " + e.getMessage());
        }
        System.out.println(result);
        return result;
    }

    private String sendToSparkAPI(String sessionId, String input) throws Exception {
        CountDownLatch latch = new CountDownLatch(1);
        AtomicReference<String> answerRef = new AtomicReference<>("");
        AtomicReference<Exception> errorRef = new AtomicReference<>();

        try {
            String authUrl = getAuthUrl(hostUrl, apiKey, apiSecret);
            Request request = new Request.Builder().url(authUrl).build();

            WebSocketListener listener = new WebSocketListener() {
                final StringBuilder currentAnswer = new StringBuilder();

                @Override
                public void onOpen(@NotNull WebSocket webSocket, @NotNull Response response) {
                    try {
                        JSONObject requestJson = buildRequestJson(input);
                        webSocket.send(requestJson.toString());
                    } catch (Exception e) {
                        errorRef.set(e);
                        latch.countDown();
                    }
                }

                @Override
                public void onMessage(@NotNull WebSocket webSocket, @NotNull String text) {
                    try {
                        JsonParse response = new Gson().fromJson(text, JsonParse.class);

                        if (response.header.code != 0) {
                            errorRef.set(new RuntimeException("API返回错误: " + response.header.code));
                            latch.countDown();
                            return;
                        }

                        // 拼接回答内容
                        for (Text t : response.payload.choices.text) {
                            currentAnswer.append(t.content);
                        }

                        // 检查是否结束
                        if (response.header.status == 2) {
                            answerRef.set(currentAnswer.toString());
                            latch.countDown();
                        }
                    } catch (Exception e) {
                        errorRef.set(e);
                        latch.countDown();
                    }
                }

                @Override
                public void onFailure(@NotNull WebSocket webSocket, @NotNull Throwable t, Response response) {
                    errorRef.set(new Exception("WebSocket连接失败: " + t.getMessage(), t));
                    latch.countDown();
                }

                @Override
                public void onClosing(@NotNull WebSocket webSocket, int code, @NotNull String reason) {
                    webSocket.close(1000, null);
                }
            };

            // 创建WebSocket连接
            CLIENT.newWebSocket(request, listener);

            // 等待响应（带超时）
            if (!latch.await(60, TimeUnit.SECONDS)) {
                throw new TimeoutException("WebSocket响应超时");
            }

            // 检查是否有错误
            if (errorRef.get() != null) {
                throw errorRef.get();
            }

            return answerRef.get();
        } finally {
            // 确保计数器释放，避免资源泄漏
            if (latch.getCount() > 0) {
                latch.countDown();
            }
        }
    }

    private JSONObject buildRequestJson(String input) {
        JSONObject requestJson = new JSONObject();

        // 构造header
        JSONObject header = new JSONObject();
        header.put("app_id", appid);
        header.put("uid", UUID.randomUUID().toString().substring(0, 10));

        // 构造parameter
        JSONObject parameter = new JSONObject();
        JSONObject chat = new JSONObject();
        chat.put("domain", domain);
        chat.put("temperature", 0.5);
        chat.put("max_tokens", 4096);
        parameter.put("chat", chat);

        // 构造payload
        JSONObject payload = new JSONObject();
        JSONObject message = new JSONObject();
        JSONArray textArray = new JSONArray();



        // 添加当前输入
        JSONObject inputObj = new JSONObject();
        inputObj.put("role", "user");
        inputObj.put("content", input);
        textArray.add(inputObj);

        message.put("text", textArray);
        payload.put("message", message);

        // 组合完整请求
        requestJson.put("header", header);
        requestJson.put("parameter", parameter);
        requestJson.put("payload", payload);

        return requestJson;
    }

    // 鉴权方法（同原始实现）
    public String getAuthUrl(String hostUrl, String apiKey, String apiSecret) throws Exception {
        URL url = new URL(hostUrl);
        // 时间
        SimpleDateFormat format = new SimpleDateFormat("EEE, dd MMM yyyy HH:mm:ss z", Locale.US);
        format.setTimeZone(TimeZone.getTimeZone("GMT"));
        String date = format.format(new Date());
        // 拼接
        String preStr = "host: " + url.getHost() + "\n" +
                "date: " + date + "\n" +
                "GET " + url.getPath() + " HTTP/1.1";
        // System.err.println(preStr);
        // SHA256加密
        Mac mac = Mac.getInstance("hmacsha256");
        SecretKeySpec spec = new SecretKeySpec(apiSecret.getBytes(StandardCharsets.UTF_8), "hmacsha256");
        mac.init(spec);

        byte[] hexDigits = mac.doFinal(preStr.getBytes(StandardCharsets.UTF_8));
        // Base64加密
        String sha = Base64.getEncoder().encodeToString(hexDigits);
        // System.err.println(sha);
        // 拼接
        String authorization = String.format("api_key=\"%s\", algorithm=\"%s\", headers=\"%s\", signature=\"%s\"", apiKey, "hmac-sha256", "host date request-line", sha);
        // 拼接地址
        HttpUrl httpUrl = Objects.requireNonNull(HttpUrl.parse("https://" + url.getHost() + url.getPath())).newBuilder().
                addQueryParameter("authorization", Base64.getEncoder().encodeToString(authorization.getBytes(StandardCharsets.UTF_8))).
                addQueryParameter("date", date).
                addQueryParameter("host", url.getHost()).
                build();

        // System.err.println(httpUrl.toString());
        return httpUrl.toString();

    }

    public boolean setModel(String model) {
        this.map  = YamlUtils.getInfo(model);
        System.out.println(map);
        return !map.isEmpty();
    }
    public boolean setFormat(Map<String, Object> format) {
        this.format = format;
        return !format.isEmpty();
    }

    public boolean setPrompt(String prompt) {
        this.prompt = prompt;
        return !prompt.isEmpty();
    }



    static class JsonParse {
        Header header;
        Payload payload;
    }

    static class Header {
        int code;
        int status;
        String sid;
    }

    static class Payload {
        Choices choices;
    }

    static class Choices {
        List<Text> text;
    }

    static class Text {
        String role;
        String content;
    }



}